{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "edb7e2c6",
   "metadata": {
    "id": "edb7e2c6"
   },
   "source": [
    "<center>\n",
    "<h4>CDS 110, Lecture 6a</h4>\n",
    "<font color=blue><h1>Trajectory Generation for a Kinematic Car Model</h1></font>\n",
    "<h3>Richard M. Murray, Winter 2024</h3>\n",
    "</center>\n",
    "\n",
    "[Open in Google Colab](https://colab.research.google.com/drive/1vBFjCU2W6fSavy8loL0JfgZyO6UC46m3)\n",
    "\n",
    "This notebook contains an example of using (optimal) trajectory generation for a vehicle steering system.  It illustrates different methods of setting up optimal control problems and solving them using python-control."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7066eb69",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import time\n",
    "try:\n",
    "  import control as ct\n",
    "  print(\"python-control\", ct.__version__)\n",
    "except ImportError:\n",
    "  !pip install control\n",
    "  import control as ct\n",
    "import control.optimal as opt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4afb09dd",
   "metadata": {
    "id": "4afb09dd"
   },
   "source": [
    "## Vehicle steering dynamics\n",
    "\n",
    "<table>\n",
    "<tr>\n",
    "    <td width=\"50%\"><img src=\"https://fbswiki.org/wiki/images/5/52/Kincar.png\" width=480></td>\n",
    "    <td width=\"50%\">\n",
    "$$\n",
    "\\large\\begin{aligned}\n",
    "  \\dot x &= \\cos\\theta\\, v \\\\\n",
    "  \\dot y &= \\sin\\theta\\, v \\\\\n",
    "  \\dot\\theta &= \\frac{v}{l} \\tan \\delta\n",
    "\\end{aligned}\n",
    "$$\n",
    "    </td>\n",
    "</tr>\n",
    "</table>\n",
    "\n",
    "The vehicle dynamics are given by a simple bicycle model.  We take the state of the system as $(x, y, \\theta)$ where $(x, y)$ is the position of the vehicle in the plane and $\\theta$ is the angle of the vehicle with respect to horizontal.  The vehicle input is given by $(v, \\delta)$ where $v$ is the forward velocity of the vehicle and $\\delta$ is the angle of the steering wheel.  The model includes saturation of the vehicle steering angle."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a6143a8a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Code to model vehicle steering dynamics\n",
    "\n",
    "# Function to compute the RHS of the system dynamics\n",
    "def kincar_update(t, x, u, params):\n",
    "    # Get the parameters for the model\n",
    "    l = params['wheelbase']             # vehicle wheelbase\n",
    "    deltamax = params['maxsteer']         # max steering angle (rad)\n",
    "\n",
    "    # Saturate the steering input\n",
    "    delta = np.clip(u[1], -deltamax, deltamax)\n",
    "\n",
    "    # Return the derivative of the state\n",
    "    return np.array([\n",
    "        np.cos(x[2]) * u[0],            # xdot = cos(theta) v\n",
    "        np.sin(x[2]) * u[0],            # ydot = sin(theta) v\n",
    "        (u[0] / l) * np.tan(delta)      # thdot = v/l tan(delta)\n",
    "    ])\n",
    "\n",
    "kincar_params={'wheelbase': 3, 'maxsteer': 0.5}\n",
    "\n",
    "# Create nonlinear input/output system\n",
    "kincar = ct.nlsys(\n",
    "    kincar_update, None, name=\"kincar\", params=kincar_params,\n",
    "    inputs=('v', 'delta'), outputs=('x', 'y', 'theta'),\n",
    "    states=('x', 'y', 'theta'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c2bf8d6-7580-4712-affc-928a8b046d8a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Utility function to plot lane change manuever\n",
    "def plot_lanechange(t, y, u, figure=None, yf=None, label=None):\n",
    "    # Plot the xy trajectory\n",
    "    plt.subplot(3, 1, 1, label='xy')\n",
    "    plt.plot(y[0], y[1], label=label)\n",
    "    plt.xlabel(\"x [m]\")\n",
    "    plt.ylabel(\"y [m]\")\n",
    "    if yf is not None:\n",
    "        plt.plot(yf[0], yf[1], 'ro')\n",
    "\n",
    "    # Plot x and y as functions of time\n",
    "    plt.subplot(3, 2, 3, label='x')\n",
    "    plt.plot(t, y[0])\n",
    "    plt.ylabel(\"$x$ [m]\")\n",
    "\n",
    "    plt.subplot(3, 2, 4, label='y')\n",
    "    plt.plot(t, y[1])\n",
    "    plt.ylabel(\"$y$ [m]\")\n",
    "\n",
    "    # Plot the inputs as a function of time\n",
    "    plt.subplot(3, 2, 5, label='v')\n",
    "    plt.plot(t, u[0])\n",
    "    plt.xlabel(\"Time $t$ [sec]\")\n",
    "    plt.ylabel(\"$v$ [m/s]\")\n",
    "\n",
    "    plt.subplot(3, 2, 6, label='delta')\n",
    "    plt.plot(t, u[1])\n",
    "    plt.xlabel(\"Time $t$ [sec]\")\n",
    "    plt.ylabel(\"$\\\\delta$ [rad]\")\n",
    "\n",
    "    plt.subplot(3, 1, 1)\n",
    "    plt.title(\"Lane change manuever\")\n",
    "    if label:\n",
    "        plt.legend()\n",
    "    plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "64bd3c3b",
   "metadata": {
    "id": "64bd3c3b"
   },
   "source": [
    "## Optimal trajectory generation\n",
    "\n",
    "The general problem we are solving is of the form:\n",
    "\n",
    "$$\n",
    "\\min_{u(\\cdot)}\n",
    "  \\int_0^T L(x,u)\\, dt + V \\bigl( x(T) \\bigr)\n",
    "$$\n",
    "subject to\n",
    "$$\n",
    "  \\dot x = f(x, u), \\qquad x\\in \\mathcal{X} \\subset \\mathbb{R}^n,\\, u\\in \\mathcal{U} \\subset \\mathbb{R}^m\n",
    "$$\n",
    "\n",
    "We consider the problem of changing from one lane to another over a perod of 10 seconds while driving at a forward speed of 10 m/s."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42dcbd79",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initial and final conditions\n",
    "x0 = np.array([  0., -2., 0.]); u0 = np.array([10., 0.])\n",
    "xf = np.array([100.,  2., 0.]); uf = np.array([10., 0.])\n",
    "Tf = 10"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ff2e044",
   "metadata": {
    "id": "5ff2e044"
   },
   "source": [
    "An important part of the optimization procedure is to give a good initial guess.  Here are some possibilities:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "650d321a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the time horizon (and spacing) for the optimization\n",
    "# timepts = np.linspace(0, Tf, 5, endpoint=True)    # Try using this and see what happens\n",
    "# timepts = np.linspace(0, Tf, 10, endpoint=True)   # Try using this and see what happens\n",
    "timepts = np.linspace(0, Tf, 20, endpoint=True)\n",
    "\n",
    "# Compute some initial guesses to use\n",
    "bend_left = [10, 0.01]          # slight left veer (will extend over all timepts)\n",
    "straight_line = (               # straight line from start to end with nominal input\n",
    "    np.array([x0 + (xf - x0) * t/Tf for t in timepts]).transpose(),\n",
    "    u0\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4e75a2c4",
   "metadata": {
    "id": "4e75a2c4"
   },
   "source": [
    "### Approach 1: standard quadratic cost\n",
    "\n",
    "We can set up the optimal control problem as trying to minimize the distance from the desired final point while at the same time as not exerting too much control effort to achieve our goal.\n",
    "\n",
    "$$\n",
    "\\min_{u(\\cdot)}\n",
    "  \\int_0^T \\left[(x(\\tau) - x_\\text{f})^T Q_x (x(\\tau) - x_\\text{f}) + (u(\\tau) - u_\\text{f})^T Q_u (u(\\tau) - u_\\text{f})\\right] \\, d\\tau\n",
    "$$\n",
    "subject to\n",
    "$$\n",
    "  \\dot x = f(x, u), \\qquad x \\in \\mathbb{R}^n,\\, u \\in \\mathbb{R}^m\n",
    "$$\n",
    "\n",
    "The optimization module solves optimal control problems by choosing the values of the input at each point in the time horizon to try to minimize the cost:\n",
    "\n",
    "$$\n",
    "u_i(t_j) = \\alpha_{i, j}, \\qquad\n",
    "u_i(t) = \\frac{t_{i+1} - t}{t_{i+1} - t_i} \\alpha_{i, j} + \\frac{t - t_i}{t_{i+1} - t_i} \\alpha_{{i+1},j}\n",
    "$$\n",
    "\n",
    "This means that each input generates a parameter value at each point in the time horizon, so the more refined your time horizon, the more parameters the optimizer has to search over."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "984c2f0b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set up the cost functions\n",
    "Qx = np.diag([.1, 10, .1])       # keep lateral error low\n",
    "Qu = np.diag([.1, 1])            # minimize applied inputs\n",
    "quad_cost = opt.quadratic_cost(kincar, Qx, Qu, x0=xf, u0=uf)\n",
    "\n",
    "# Compute the optimal control, setting step size for gradient calculation (eps)\n",
    "start_time = time.process_time()\n",
    "result1 = opt.solve_ocp(\n",
    "    kincar, timepts, x0, quad_cost,\n",
    "    initial_guess=straight_line,\n",
    "    # initial_guess= bend_left,\n",
    "    # initial_guess=u0,\n",
    "    # minimize_method='trust-constr',\n",
    "    # minimize_options={'finite_diff_rel_step': 0.01},\n",
    "    # trajectory_method='shooting'\n",
    "    # solve_ivp_method='LSODA'\n",
    ")\n",
    "print(\"* Total time = %5g seconds\\n\" % (time.process_time() - start_time))\n",
    "\n",
    "# Plot the results from the optimization\n",
    "plot_lanechange(timepts, result1.states, result1.inputs, xf)\n",
    "print(\"Final computed state: \", result1.states[:,-1])\n",
    "\n",
    "# Simulate the system and see what happens\n",
    "t1, u1 = result1.time, result1.inputs\n",
    "t1, y1 = ct.input_output_response(kincar, timepts, u1, x0)\n",
    "plot_lanechange(t1, y1, u1, yf=xf[0:2])\n",
    "print(\"Final simulated state:\", y1[:,-1])\n",
    "\n",
    "# Label the different lines\n",
    "plt.subplot(3, 1, 1)\n",
    "plt.legend(['desired', 'simulated', 'endpoint'])\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7cade52",
   "metadata": {
    "id": "b7cade52"
   },
   "source": [
    "Note the amount of time required to solve the problem and also any warning messages about to being able to solve the optimization (mainly in earlier versions of python-control).  You can try to adjust a number of factors to try to get a better solution:\n",
    "* Try changing the number of points in the time horizon\n",
    "* Try using a different initial guess\n",
    "* Try changing the optimization method (see commented out code)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a9f9d9b",
   "metadata": {
    "id": "6a9f9d9b"
   },
   "source": [
    "### Approach 2: input cost, input constraints, terminal cost\n",
    "\n",
    "The previous solution integrates the position error for the entire horizon, and so the car changes lanes very quickly (at the cost of larger inputs).  Instead, we can penalize the final state and impose a higher cost on the inputs, resulting in a more gradual lane change.\n",
    "\n",
    "$$\n",
    "\\min_{u(\\cdot)}\n",
    "  \\int_0^T \\underbrace{\\left[x(\\tau)^T Q_x x(\\tau) + (u(\\tau) - u_\\text{f})^T Q_u (u(\\tau) - u_\\text{f})\\right]}_{L(x, u)} \\, d\\tau + \\underbrace{(x(T) - x_\\text{f})^T Q_\\text{f} (x(T) - x_\\text{f})}_{V\\left(x(T)\\right)}\n",
    "$$\n",
    "subject to\n",
    "$$\n",
    "  \\dot x = f(x, u), \\qquad x \\in \\mathbb{R}^n,\\, u \\in \\mathbb{R}^m\n",
    "$$\n",
    "\n",
    "We can also try using a different solver for this example.  You can pass the solver using the `minimize_method` keyword and send options to the solver using the `minimize_options` keyword (which should be set to a dictionary of options)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a201e33c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add input constraint, input cost, terminal cost\n",
    "constraints = [ opt.input_range_constraint(kincar, [8, -0.1], [12, 0.1]) ]\n",
    "traj_cost = opt.quadratic_cost(kincar, None, np.diag([0.1, 1]), u0=uf)\n",
    "term_cost = opt.quadratic_cost(kincar, np.diag([1, 10, 100]), None, x0=xf)\n",
    "\n",
    "# Compute the optimal control\n",
    "start_time = time.process_time()\n",
    "result2 = opt.solve_ocp(\n",
    "    kincar, timepts, x0, traj_cost, constraints, terminal_cost=term_cost,\n",
    "    initial_guess=straight_line,\n",
    "    # minimize_method='trust-constr',\n",
    "    # minimize_options={'finite_diff_rel_step': 0.01},\n",
    "    # minimize_method='SLSQP', minimize_options={'eps': 0.01},\n",
    "    # log=True,\n",
    ")\n",
    "print(\"* Total time = %5g seconds\\n\" % (time.process_time() - start_time))\n",
    "\n",
    "# Plot the results from the optimization\n",
    "plot_lanechange(timepts, result2.states, result2.inputs, xf)\n",
    "print(\"Final computed state: \", result2.states[:,-1])\n",
    "\n",
    "# Simulate the system and see what happens\n",
    "t2, u2 = result2.time, result2.inputs\n",
    "t2, y2 = ct.input_output_response(kincar, timepts, u2, x0)\n",
    "plot_lanechange(t2, y2, u2, yf=xf[0:2])\n",
    "print(\"Final simulated state:\", y2[:,-1])\n",
    "\n",
    "# Label the different lines\n",
    "plt.subplot(3, 1, 1)\n",
    "plt.legend(['desired', 'simulated', 'endpoint'], loc='upper left')\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3d2ccf97",
   "metadata": {
    "id": "3d2ccf97"
   },
   "source": [
    "### Approach 3: terminal constraints\n",
    "\n",
    "We can also remove the cost function on the state and replace it with a terminal *constraint* on the state as well as bounds on the inputs.  If a solution is found, it guarantees we get to exactly the final state:\n",
    "\n",
    "$$\n",
    "\\min_{u(\\cdot)}\n",
    "  \\int_0^T \\underbrace{(u(\\tau) - u_\\text{f})^T Q_u (u(\\tau) - u_\\text{f})}_{L(x, u)} \\, d\\tau\n",
    "$$\n",
    "subject to\n",
    "$$\n",
    "  \\begin{aligned}\n",
    "  \\dot x &= f(x, u), & \\qquad &x \\in \\mathbb{R}^n,\\, u \\in \\mathbb{R}^m \\\\\n",
    "  x(T) &= x_\\text{f} & &u_\\text{lb} \\leq u(t) \\leq u_\\text{ub},\\, \\text{for all $t$}\n",
    "  \\end{aligned}\n",
    "$$\n",
    "\n",
    "Note that trajectory and terminal constraints can be very difficult to satisfy for a general optimization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc77a856",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Input cost and terminal constraints\n",
    "R = np.diag([1, 1])                 # minimize applied inputs\n",
    "cost3 = opt.quadratic_cost(kincar, np.zeros((3,3)), R, u0=uf)\n",
    "constraints = [\n",
    "    opt.input_range_constraint(kincar, [8, -0.1], [12, 0.1]) ]\n",
    "terminal = [ opt.state_range_constraint(kincar, xf, xf) ]\n",
    "\n",
    "# Compute the optimal control\n",
    "start_time = time.process_time()\n",
    "result3 = opt.solve_ocp(\n",
    "    kincar, timepts, x0, cost3, constraints,\n",
    "    terminal_constraints=terminal, initial_guess=straight_line,\n",
    "#    solve_ivp_kwargs={'atol': 1e-3, 'rtol': 1e-2},\n",
    "#    minimize_method='trust-constr',\n",
    "#    minimize_options={'finite_diff_rel_step': 0.01},\n",
    ")\n",
    "print(\"* Total time = %5g seconds\\n\" % (time.process_time() - start_time))\n",
    "\n",
    "# Plot the results from the optimization\n",
    "plot_lanechange(timepts, result3.states, result3.inputs, xf)\n",
    "print(\"Final computed state: \", result3.states[:,-1])\n",
    "\n",
    "# Simulate the system and see what happens\n",
    "t3, u3 = result3.time, result3.inputs\n",
    "t3, y3 = ct.input_output_response(kincar, timepts, u3, x0)\n",
    "plot_lanechange(t3, y3, u3, yf=xf[0:2])\n",
    "print(\"Final state: \", y3[:,-1])\n",
    "\n",
    "# Label the different lines\n",
    "plt.subplot(3, 1, 1)\n",
    "plt.legend(['desired', 'simulated', 'endpoint'], loc='upper left')\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9e744463",
   "metadata": {
    "id": "9e744463"
   },
   "source": [
    "### Approach 4: terminal constraints w/ basis functions (if time)\n",
    "\n",
    "As a final example, we can use a basis function to reduce the size of the problem and get faster answers with more temporal resolution:\n",
    "\n",
    "$$\n",
    "\\min_{u(\\cdot)}\n",
    "  \\int_0^T L(x, u) \\, d\\tau + V\\left(x(T)\\right)\n",
    "$$\n",
    "subject to\n",
    "$$\n",
    "  \\begin{aligned}\n",
    "  \\dot x &= f(x, u), \\qquad x \\in \\mathcal{X} \\subset \\mathbb{R}^n,\\, u \\in \\mathcal{U} \\subset \\mathbb{R}^m \\\\\n",
    "  u(t) &= \\sum_i \\alpha_i \\phi^i(t),\n",
    "  \\end{aligned}\n",
    "$$\n",
    "where $\\phi^i(t)$ are a set of basis functions.\n",
    "\n",
    "Here we parameterize the input by a set of 4 Bezier curves but solve for a much more time resolved set of inputs.  Note that while we are using the `control.flatsys` module to define the basis functions, we are not exploiting the fact that the system is differentially flat."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ee82aa25",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get basis functions for flat systems module\n",
    "import control.flatsys as flat\n",
    "\n",
    "# Compute the optimal control\n",
    "start_time = time.process_time()\n",
    "result4 = opt.solve_ocp(\n",
    "    kincar, timepts, x0, quad_cost, constraints,\n",
    "    terminal_constraints=terminal,\n",
    "    initial_guess=straight_line,\n",
    "    basis=flat.PolyFamily(4, T=Tf),\n",
    "    # solve_ivp_kwargs={'method': 'RK45', 'atol': 1e-2, 'rtol': 1e-2},\n",
    "    # solve_ivp_kwargs={'atol': 1e-3, 'rtol': 1e-2},\n",
    "    # minimize_method='trust-constr', minimize_options={'disp': True},\n",
    "    log=False\n",
    ")\n",
    "print(\"* Total time = %5g seconds\\n\" % (time.process_time() - start_time))\n",
    "\n",
    "# Plot the results from the optimization\n",
    "plot_lanechange(timepts, result4.states, result4.inputs, xf)\n",
    "print(\"Final computed state: \", result3.states[:,-1])\n",
    "\n",
    "# Simulate the system and see what happens\n",
    "t4, u4 = result4.time, result4.inputs\n",
    "t4, y4 = ct.input_output_response(kincar, timepts, u4, x0)\n",
    "plot_lanechange(t4, y4, u4, yf=xf[0:2])\n",
    "print(\"Final simulated state: \", y4[:,-1])\n",
    "\n",
    "# Label the different lines\n",
    "plt.subplot(3, 1, 1)\n",
    "plt.legend(['desired', 'simulated', 'endpoint'], loc='upper left')\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a74388e",
   "metadata": {
    "id": "2a74388e"
   },
   "source": [
    "Note how much smoother the inputs look, although the solver can still have a hard time satisfying the final constraints, resulting in longer computation times."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1465d149",
   "metadata": {
    "id": "1465d149"
   },
   "source": [
    "### Additional things to try\n",
    "\n",
    "* Compare the results here with what we go last week exploiting the property of differential flatness (computation time, in particular)\n",
    "* Try using different weights, solvers, initial guess and other properties and see how things change.\n",
    "* Try using different values for `initial_guess` to get faster convergence and/or different classes of solutions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02bad3d5",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
